from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from contextlib import asynccontextmanager
import asyncio
import json
import os
import logging

# 导入新的Connectors
from utils.connectors.internal_auth import AuthManager
from utils.connectors.siem_connector import SiemConnector

# 导入PocketFlow流程创建函数
from flow import create_soc_analysis_flow

# --- 全局资源初始化 ---
# 在应用级别创建单例的AuthManager，可以被所有会话共享
# 这样可以确保Token在整个应用生命周期内尽可能地被复用
try:
    auth_manager = AuthManager()
    print("[Global] AuthManager 初始化成功。")
except ValueError as e:
    print(f"[CRITICAL] 全局AuthManager初始化失败: {e}")
    print("请立即检查 .env 文件中的 INTERNAL_API_* 配置项。")
    auth_manager = None

# --- FastAPI 应用初始化 --- #
active_connections: dict[str, WebSocket] = {}
running_tasks: dict[str, dict] = {}

@asynccontextmanager
async def lifespan(app: FastAPI):
    print("FastAPI 应用启动...")
    yield
    print("FastAPI 应用关闭。")

app = FastAPI(lifespan=lifespan)

# 挂载前端静态文件
# 前端文件在当前项目的frontend目录下
frontend_path = os.path.abspath(os.path.join(os.path.dirname(__file__), 'frontend'))

# 检查前端路径是否存在
if not os.path.exists(frontend_path):
    print(f"警告: 未找到前端静态文件目录: {frontend_path}")
    print("请确保frontend文件夹存在于项目根目录。")
else:
    print(f"前端静态文件目录: {frontend_path}")

app.mount("/static", StaticFiles(directory=frontend_path), name="static")

@app.get("/")
async def get_index():
    # 返回前端的index.html
    index_html_path = os.path.join(frontend_path, 'index.html')
    if not os.path.exists(index_html_path):
        return HTMLResponse(content="<h1>Error: index.html not found!</h1><p>Please ensure the frontend files are correctly placed in the 'frontend' directory.</p>", status_code=404)
    with open(index_html_path, 'r', encoding='utf-8') as f:
        return HTMLResponse(content=f.read())

@app.get("/c_hello")
async def handle_c_hello(asker: str = None):
    """
    处理神秘的c_hello请求，可能是前端的健康检查或监控请求
    """
    return {"status": "ok", "message": "SOC Analyzer is running", "asker": asker}

@app.get("/api/status")
async def get_system_status():
    """
    获取系统状态，包括活跃连接和运行任务数量
    """
    return {
        "status": "running",
        "active_connections": len(active_connections),
        "running_tasks": len(running_tasks),
        "task_details": {
            client_id: {
                "analysis": "running" if not tasks["analysis"].done() else "completed",
                "queue_consumer": "running" if not tasks["queue_consumer"].done() else "completed"
            } if isinstance(tasks, dict) else (
                "running" if not tasks.done() else "completed"
            )
            for client_id, tasks in running_tasks.items()
        }
    }

# --- WebSocket 端点 --- #
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    client_id = str(websocket.client.host) + ":" + str(websocket.client.port)
    active_connections[client_id] = websocket
    print(f"WebSocket连接建立: {client_id}")

    try:
        while True:
            # 接收前端发送的告警数据
            data = await websocket.receive_text()
            message = json.loads(data)
            
            if message.get("type") == "start_analysis":
                # 检查是否已有任务在运行
                if client_id in running_tasks:
                    tasks = running_tasks[client_id]
                    if isinstance(tasks, dict):
                        # 新的任务结构
                        if not tasks["analysis"].done():
                            await websocket.send_json({
                                "type": "error",
                                "step": "websocket",
                                "data": {"message": "已有分析任务正在运行，请等待完成后再提交新任务"}
                            })
                            continue
                    else:
                        # 旧的任务结构（向后兼容）
                        if not tasks.done():
                            await websocket.send_json({
                                "type": "error",
                                "step": "websocket",
                                "data": {"message": "已有分析任务正在运行，请等待完成后再提交新任务"}
                            })
                            continue

                raw_alert_json_str = message.get("alert_data")
                if not raw_alert_json_str:
                    await websocket.send_json({"type": "error", "step": "websocket", "data": {"message": "未接收到告警数据"}})
                    continue

                # --- 注入Connectors和依赖 --- #
                if not auth_manager:
                    await websocket.send_json({"type": "error", "step": "system", "data": {"message": "系统错误：内部认证服务未初始化，无法开始分析。"}})
                    continue
                
                # 为本次会话创建专用的Connector实例，它们共享全局的AuthManager
                siem_connector = SiemConnector(auth_manager)
                # TODO: 在此可以添加 CmdbConnector 等其他连接器

                # 将状态队列、原始告警和Connectors实例传递给PocketFlow
                shared_data = {
                    "raw_alert": raw_alert_json_str,
                    "status_queue": asyncio.Queue(),
                    "websocket": websocket,
                    "connectors": {
                        "siem": siem_connector,
                        # "cmdb": cmdb_connector # 示例
                    }
                }

                # 启动两个后台任务：
                # 1. 运行PocketFlow分析流程
                # 2. 消费状态队列并转发到前端
                analysis_task = asyncio.create_task(run_pocketflow_analysis(shared_data, client_id))
                queue_consumer_task = asyncio.create_task(consume_status_queue(shared_data["status_queue"], client_id))

                # 将两个任务都存储起来，以便后续清理
                running_tasks[client_id] = {
                    "analysis": analysis_task,
                    "queue_consumer": queue_consumer_task
                }

    except WebSocketDisconnect:
        # 清理连接和任务
        if client_id in active_connections:
            del active_connections[client_id]
        if client_id in running_tasks:
            tasks = running_tasks[client_id]
            if isinstance(tasks, dict):
                # 新的任务结构，取消所有任务
                tasks["analysis"].cancel()
                tasks["queue_consumer"].cancel()
            else:
                # 旧的任务结构（向后兼容）
                tasks.cancel()
            del running_tasks[client_id]
        print(f"WebSocket连接断开: {client_id}")
    except Exception as e:
        print(f"WebSocket错误: {e}")
        if client_id in active_connections:
            await active_connections[client_id].send_json({"type": "error", "step": "websocket", "data": {"message": str(e)}})

async def consume_status_queue(status_queue: asyncio.Queue, client_id: str):
    """
    消费状态队列中的消息，并将其转发到前端WebSocket。
    """
    try:
        while True:
            # 等待队列中的消息
            message = await status_queue.get()

            # 如果收到None，表示队列结束
            if message is None:
                print(f"状态队列消费者收到结束信号: {client_id}")
                break

            # 将消息转发到前端
            if client_id in active_connections:
                await active_connections[client_id].send_json(message)
                print(f"转发状态消息到前端: {message.get('type', 'unknown')} - {message.get('step', 'unknown')}")

            # 标记任务完成
            status_queue.task_done()

    except asyncio.CancelledError:
        print(f"状态队列消费者被取消: {client_id}")
    except Exception as e:
        print(f"状态队列消费者错误: {e}")
        if client_id in active_connections:
            await active_connections[client_id].send_json({
                "type": "error",
                "step": "queue_consumer",
                "data": {"message": f"状态队列消费者错误: {str(e)}"}
            })

async def run_pocketflow_analysis(shared_data: dict, client_id: str):
    """
    在后台运行PocketFlow分析流程，并将状态推送到前端。
    """
    soc_flow = create_soc_analysis_flow()
    full_report_content = ""

    # 从 shared_data 中获取 status_queue
    status_queue = shared_data.get("status_queue")
    if not status_queue:
        print(f"[ERROR] run_pocketflow_analysis: status_queue not found in shared_data for client {client_id}")
        if client_id in active_connections:
            await active_connections[client_id].send_json({
                "type": "error",
                "step": "analysis_error",
                "data": {"message": "系统错误：状态队列未找到，无法进行分析。"}
            })
        return

    try:
        # PocketFlow的run_async现在会通过shared_data["status_queue"]发送所有状态更新
        # 并且GenerateReportNode会直接通过shared_data["websocket"]发送流式报告chunk
        # 所以这里不再需要迭代soc_flow.run_async的返回值
        await soc_flow.run_async(shared_data)

        # 流程结束后，从shared_data中获取最终报告
        final_report_from_shared = shared_data.get("final_report", "流程未生成报告。")
        
        # 发送最终报告完成信号
        if client_id in active_connections:
            await active_connections[client_id].send_json({
                "type": "final_report_complete",
                "step": "GenerateReportNode",
                "data": {"report_content": final_report_from_shared}
            })

    except Exception as e:
        print(f"PocketFlow分析错误: {e}")
        if client_id in active_connections:
            await active_connections[client_id].send_json({
                "type": "error",
                "step": "analysis_error",
                "data": {"message": f"分析过程中发生错误: {str(e)}"}
            })
    finally:
        # 确保队列被清空，并发送结束信号
        while not status_queue.empty():
            await status_queue.get() # 清空队列
        await status_queue.put(None) # 发送结束信号给可能的消费者 (虽然这里没有直接消费者)

        # 清理任务状态
        if client_id in running_tasks:
            tasks = running_tasks[client_id]
            if isinstance(tasks, dict):
                # 确保队列消费者也被取消
                if not tasks["queue_consumer"].done():
                    tasks["queue_consumer"].cancel()
            del running_tasks[client_id]

        print(f"PocketFlow分析任务完成或出错: {client_id}")

        # 发送任务完成信号
        if client_id in active_connections:
            await active_connections[client_id].send_json({
                "type": "analysis_complete",
                "step": "system",
                "data": {"message": "分析任务已完成"}
            })

# --- 运行FastAPI应用 --- #
# 在命令行中运行: uvicorn server:app --reload --port 8000
